import matplotlib
import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F


plt.rcParams['figure.figsize'] = (7, 3.5)
plt.rcParams['figure.dpi'] = 150
plt.rcParams['axes.unicode_minus'] = False  #解决坐标轴负数的铅显示问题

"""
relu(x) = max(x, 0)

- 优点：
    ReLU激活函数是一个简单的计算，如果输入大于0，直接返回作为输入提供的值；如果输入是0或更小，返回值0。

    - 相较于sigmoid函数以及Tanh函数来看，在输入为正时，Relu函数不存在饱和问题，即解决了gradient vanishing问题，使得深层网络可训练
    - Relu输出会使一部分神经元为0值，在带来网络稀疏性的同时，也减少了参数之间的关联性，一定程度上缓解了过拟合的问题
    - 计算速度非常快
    - 收敛速度远快于sigmoid以及Tanh函数

- 缺点：
    - 输出不是zero-centered
    - 存在Dead Relu Problem，即某些神经元可能永远不会被激活，进而导致相应参数一直得不到更新，产生该问题主要原因包括参数初始化问题以及学习率设置过大问题
    - ReLU不会对数据做幅度压缩，所以数据的幅度会随着模型层数的增加不断扩张，当输入为正值，导数为1，在“链式反应”中，不会出现梯度消失，但梯度下降的强度则完全取决于权值的乘积，如此可能会导致梯度爆炸问题
"""

#! 1. 定义可以进行前项和后算运算ReLU算子
class MyReLUAutoGrad(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x):
        ctx.save_for_backward(x)       # 保存上下文用于反向传播
        return F.relu(x)
    
    @staticmethod
    def backward(ctx, grad_output):   # y.backward(torch.ones_like(y)) 传进来的
        x, = ctx.saved_tensors
        grad = grad_output * torch.where(x < 0., torch.zeros_like(x), torch.ones_like(x))
        return grad
    

class ReLU(nn.Module):
    def __init__(self):
        super(ReLU, self).__init__()
    
    def forward(self, x):
        return MyReLUAutoGrad.apply(x)
    
#! 2. 测试
x = torch.randn(2, 10, dtype=torch.float32, requires_grad=True)
print(x)

net = ReLU()
y   = net(x)
""" 在 PyTorch 中，loss.backward(retain_graph=True) 的 retain_graph=True 参数用于指示在执行反向传播后保留计算图。
这个参数在某些特定的情况下非常有用，尤其是当您需要多次对同一个图进行反向传播时。
"""
y.backward(torch.ones_like(y), retain_graph=True)        # 会调用到MyReLUAutoGrad.backward
# tensor([[1., 1., 0., 1., 0., 1., 1., 0., 0., 1.],
#         [1., 0., 0., 0., 0., 1., 0., 0., 1., 1.]])
print(x.grad)

# tensor([[2., 2., 0., 2., 0., 2., 2., 0., 0., 2.],
#         [2., 0., 0., 0., 0., 2., 0., 0., 2., 2.]])
y.backward(torch.ones_like(y), retain_graph=True)    
print(x.grad)

print(x.grad.zero_())              #! 梯度清零
y.backward(torch.ones_like(y), retain_graph=True)   
# tensor([[1., 1., 0., 1., 0., 1., 1., 0., 0., 1.],
#         [1., 0., 0., 0., 0., 1., 0., 0., 1., 1.]])
print(x.grad)

#! 3. 梯度可视化
# visualization
# inp = torch.arange(-8, 8, 0.05, requires_grad=True)
# out = net(inp)
# out.sum().backward()

# inp_grad = inp.grad

# plt.plot(inp.detach().numpy(),
#          out.detach().numpy(),
#          label=r"$relu(x)$",
#          alpha=0.7)
# plt.plot(inp.detach().numpy(),
#          inp_grad.numpy(),
#          label=r"$relu'(x)$",
#          alpha=0.5)
# plt.scatter(0, 0, color='None', marker='o', edgecolors='r', s=50)
# plt.grid()
# plt.legend()
# plt.show()
